package com.vrv.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.vrv.entity.ElasticsearchProperties;
import com.vrv.entity.EsProp;
import com.vrv.entity.Result;
import com.vrv.utils.EsUtils;
import io.swagger.annotations.ApiImplicitParam;
import io.swagger.annotations.ApiParam;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.admin.indices.get.GetIndexResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.time.Duration;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;

/**
 * @author vrv
 * @ClassName: EsController
 * @Desc: TODO
 * @Date: 2024/07/04 16:06
 * @Version: 1.0
 * @Modified By:
 */
@RestController
@RequestMapping("es")
@Slf4j
public class EsController {

    /**
     * 0:进行中
     * 1:完成
     */
    private static final Cache<EsProp, Integer> CACHE = CacheBuilder.newBuilder().maximumSize(1000).build();

    private static final Cache<EsProp, Set<String>> INDEX_CACHE = CacheBuilder.newBuilder().maximumSize(1000).expireAfterWrite(Duration.ofHours(24)).build();

    @PostMapping("sync")
    public Object sync(@ApiParam(required = false, value = "0:断点 1:强制", defaultValue = "1") @RequestParam(defaultValue = "1", required = false) String force, @RequestBody EsProp prop) {
        Integer ifPresent = CACHE.getIfPresent(prop);
        if (null != ifPresent && 0 == ifPresent) {
            return Result.<String>builder().msg("任务运行中").code(200).build();
        }

        EsProp.Prop source = prop.getSource();
        RestHighLevelClient sourceClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(source.getUrls(), source.getUsername(), source.getPassword()));
        EsProp.Prop target = prop.getTarget();
        RestHighLevelClient targetClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(target.getUrls(), target.getUsername(), target.getPassword()));

        try {
            CACHE.put(prop, 0);
            syncIndexTemplate(sourceClient, targetClient);
            syncIndex("1".equals(force), prop, sourceClient, targetClient);
            // 任务已经完成
            CACHE.put(prop, 1);
        } catch (Exception e) {
            // 任务失败
            CACHE.put(prop, 2);
            return Result.<String>builder().msg("任务失败").data(e.getMessage()).code(500).build();
        } finally {
            try {
                sourceClient.close();
            } catch (IOException e) {
                log.error("sourceClient error={}", e.getMessage());
            }
            try {
                targetClient.close();
            } catch (IOException e) {
                log.error("targetClient error={}", e.getMessage());
            }
        }
        return Result.<String>builder().msg("任务成功").code(200).build();
    }


    @PostMapping("syncTemplate")
    public Object syncTemplate(@RequestBody EsProp prop) {
        EsProp.Prop source = prop.getSource();
        RestHighLevelClient sourceClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(source.getUrls(), source.getUsername(), source.getPassword()));
        EsProp.Prop target = prop.getTarget();
        RestHighLevelClient targetClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(target.getUrls(), target.getUsername(), target.getPassword()));
        try {
            syncIndexTemplate(sourceClient, targetClient);
            // 任务已经完成
        } catch (Exception e) {
            // 任务失败
            return Result.<String>builder().msg("任务失败").data(e.getMessage()).code(500).build();
        } finally {
            try {
                sourceClient.close();
            } catch (IOException e) {
                log.error("sourceClient error={}", e.getMessage());
            }
            try {
                targetClient.close();
            } catch (IOException e) {
                log.error("targetClient error={}", e.getMessage());
            }
        }
        return Result.<String>builder().msg("任务成功").code(200).build();
    }


    @PostMapping("syncData")
    public Object syncData(@ApiParam(value = "0:断点 1:强制", defaultValue = "1") @RequestParam(defaultValue = "1", required = false) String force, @RequestBody EsProp prop) {
        EsProp.Prop source = prop.getSource();
        RestHighLevelClient sourceClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(source.getUrls(), source.getUsername(), source.getPassword()));
        EsProp.Prop target = prop.getTarget();
        RestHighLevelClient targetClient = EsUtils.getRestHighLevelClient(new ElasticsearchProperties(target.getUrls(), target.getUsername(), target.getPassword()));
        try {
            syncIndex("1".equals(force), prop, sourceClient, targetClient);
            // 任务已经完成
        } catch (Exception e) {
            // 任务失败
            return Result.<String>builder().msg("任务失败").data(e.getMessage()).code(500).build();
        } finally {
            try {
                sourceClient.close();
            } catch (IOException e) {
                log.error("sourceClient error={}", e.getMessage());
            }
            try {
                targetClient.close();
            } catch (IOException e) {
                log.error("targetClient error={}", e.getMessage());
            }
        }
        return Result.<String>builder().msg("任务成功").code(200).build();
    }


    public int syncIndex(boolean force, EsProp prop, RestHighLevelClient sourceClient, RestHighLevelClient targetClient) throws IOException, ExecutionException {
        log.info("Starting index synchronization...");
        Set<String> set = INDEX_CACHE.get(prop, new Callable<Set<String>>() {
            @Override
            public Set<String> call() throws Exception {
                return new HashSet<>();
            }
        });

        // 获取所有索引名称
        GetIndexRequest getIndexRequest = new GetIndexRequest().indices("*");
        GetIndexResponse getIndexResponse = sourceClient.indices().get(getIndexRequest, RequestOptions.DEFAULT);
        String[] indices = getIndexResponse.getIndices();

        log.info("Found {} indices to synchronize.", indices.length);

        // 遍历所有索引并进行同步
        for (String index : indices) {
            if (index.startsWith(".")) {
                log.info("skip Synchronizing index: {}", index);
            } else {
                log.info("Synchronizing index: {}", index);
                try {
                    if (!force && set.contains(index)) {
                        // 如果已经传输过了 则跳过当前索引
                        continue;
                    }
                    syncSingleIndex(sourceClient, targetClient, index);
                    // 传输完成后将该索引添加进去
                    set.add(index);
                } catch (Exception e) {
                    e.printStackTrace();
                    log.error("同步索引数据失败 {} error", index, e.getMessage());
                }
            }
        }

        log.info("Index synchronization completed.");
        return indices.length;
    }

    private void syncSingleIndex(RestHighLevelClient sourceClient, RestHighLevelClient targetClient, String index) throws IOException {
        if (!EsUtils.exist(sourceClient, index)) {
            log.error("索引不存在 {}", index);
            return;
        }
        String scrollId = null;
        final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(1L));
        SearchRequest searchRequest = new SearchRequest(index);
        searchRequest.scroll(scroll);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.sort("_doc", SortOrder.ASC); // Optimize for scroll context
        searchSourceBuilder.size(1000); // Fetch 1000 documents per scroll
        searchRequest.source(searchSourceBuilder);

        // Initial search to get the first scroll ID
        SearchResponse searchResponse = sourceClient.search(searchRequest, RequestOptions.DEFAULT);
        scrollId = searchResponse.getScrollId();

        // Loop until no more hits are returned
        while (true) {
            BulkRequest bulkRequest = new BulkRequest();
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                // Create index request for each document
                IndexRequest indexRequest = new IndexRequest(hit.getIndex());
                indexRequest.id(hit.getId());
                indexRequest.source(hit.getSourceAsMap());
                bulkRequest.add(indexRequest);
            }

            if (bulkRequest.numberOfActions() > 0) {
                // Execute bulk request
                BulkResponse bulkResponse = targetClient.bulk(bulkRequest, RequestOptions.DEFAULT);
                if (bulkResponse.hasFailures()) {
                    // Handle the failures (log them, reprocess them, etc.)
                    log.error("Bulk request had failures: {}", bulkResponse.buildFailureMessage());
                } else {
                    log.info("Successfully synchronized {} documents for index {}", bulkRequest.numberOfActions(), index);
                }
            }

            // Prepare next scroll request
            SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
            scrollRequest.scroll(scroll);
            searchResponse = sourceClient.scroll(scrollRequest, RequestOptions.DEFAULT);

            // Break if no more hits are returned
            if (searchResponse.getHits().getHits().length == 0) {
                break;
            }
        }

        // Clear the scroll context
        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
        clearScrollRequest.addScrollId(scrollId);
        sourceClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);

        log.info("Finished synchronizing index: {}", index);
    }

    private int syncIndexTemplate(RestHighLevelClient sourceClient, RestHighLevelClient targetClient) throws IOException {
        String allTemplate = EsUtils.getAllTemplate(sourceClient);

        if (StringUtils.hasText(allTemplate)) {
            JSONObject jsonObject = JSON.parseObject(allTemplate);
            int size = jsonObject.keySet().size();
            log.info("syncIndexTemplate size={}", size);
            for (String key : jsonObject.keySet()) {
                log.info("start sync IndexTemplate name={}", key);
                EsUtils.putTemplate(targetClient, key, jsonObject.getString(key));
            }
            log.info("Finished syncIndexTemplate ");
            return size;
        }
        return 0;
    }
}
